import pysam
import argparse
import pod5 as p5
from tqdm import tqdm
from pathlib import Path
from collections import defaultdict
import gzip

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--inbam', type=str, default="/public/data/wangjianxin/nipeng/workspace/arab_col0_csu2_nextomics.ont_r10.4.1_simplex.pod5_pass.bam", required=False,
                        help="bam ready to be split")
    parser.add_argument('--outbam', type=str, default="/public/home/hpc234701005/data/R10.4/arab/train/arab_col0_train.bam", required=False,
                        help="where the splited bam to store")
    parser.add_argument('--pod5_dir', type=str, required=False,
                        help="directory containing POD5 files")
    parser.add_argument('--tsv', type=str, required=False,
                        help="tsv file with read IDs (can be .gz compressed)")
    parser.add_argument('--stat', action='store_true', default=False, help="stat mode")
    return parser.parse_args()

def read_id(pod5_dir):
    """创建read_id到POD5文件名的映射，同时返回read_id集合"""
    readid_to_pod5 = defaultdict(list)
    recursive = True
    glob = Path.rglob if recursive else Path.glob
    
    for pod5_path in tqdm(glob(Path(pod5_dir), "*.pod5"), desc="read ids"):
        with p5.DatasetReader(pod5_path, recursive=True, max_cached_readers=1) as dataset:
            for read_record in dataset:
                read_id = str(read_record.read_id)
                pod5_filename = pod5_path.name
                readid_to_pod5[read_id].append(pod5_filename)
    
    # 将多重映射简化为单一映射（取第一个文件名），并创建ids集合
    ids_set = set(readid_to_pod5.keys())
    final_mapping = {k: v[0] for k, v in readid_to_pod5.items()}
    print(f"Found {len(ids_set)} unique read IDs in POD5 files")
    return ids_set, final_mapping

def read_tsv(key_input):
    """读取TSV文件中的read IDs，支持.gz压缩文件"""
    key_indexes = set()
    opener = gzip.open if key_input.endswith('.gz') else open
    mode = 'rt' if key_input.endswith('.gz') else 'r'
    with opener(key_input, mode) as input_file:
        for line in input_file:
            key = line.strip()
            key_indexes.add(key)
    return key_indexes, {}  # 返回空映射，因为TSV不提供文件名

def calculate_median(array):
    sorted_array = sorted(array)
    length = len(sorted_array)
    if length % 2 == 0:
        middle_left = sorted_array[length // 2 - 1]
        middle_right = sorted_array[length // 2]
        median = (middle_left + middle_right) / 2
    else:
        median = sorted_array[length // 2]
    return median

def calculate_mean(array):
    total = sum(array)
    mean = total / len(array)
    return mean

def deal_bam(inbam, outbam, ids, pod5_mapping, stat=False):
    input_bam = pysam.AlignmentFile(inbam, 'rb', check_sq=False)
    if not stat:
        output_bam = pysam.AlignmentFile(outbam, 'wb', header=input_bam.header)
    
    read_num = 0
    query_length = []
    modified_fn_count = 0
    
    for read in tqdm(input_bam, desc="deal bam"):
        read_id = read.query_name
        pi_tag = read.get_tag("pi") if read.has_tag("pi") else None
        
        # 检查是否匹配（基于query_name或pi）
        if read_id in ids or (pi_tag and pi_tag in ids):
            read_num += 1
            query_length.append(read.query_length)
            
            # 检查并更新fn标签
            target_id = read_id if read_id in pod5_mapping else pi_tag
            if target_id in pod5_mapping:
                read.set_tag('fn', pod5_mapping[target_id], value_type='Z')
                modified_fn_count += 1
            
            if not stat:
                output_bam.write(read)
    
    print(f'basecalled successful {read_num} read in {len(ids)} read, proportion is {read_num/len(ids):.2%}')
    if query_length:
        print(f'the basecalled successful read\'s mean is {calculate_mean(query_length):.2f} and median is {calculate_median(query_length):.2f}')
    print(f'Modified fn tag for {modified_fn_count} reads')
    
    input_bam.close()
    if not stat:
        output_bam.close()

if __name__ == '__main__':
    args = parse_args()
    if args.tsv is not None:
        print(f"Reading TSV: {args.tsv}")
        ids, pod5_mapping = read_tsv(args.tsv)
    else:
        print(f"Processing POD5 directory: {args.pod5_dir}")
        ids, pod5_mapping = read_id(args.pod5_dir)
    
    print('start deal bam')
    deal_bam(args.inbam, args.outbam, ids, pod5_mapping, args.stat)